import matplotlib.pyplot as plt
import numpy as np


# 从txt文件读取数据
def read_data(filename):
    data = []
    with open(filename, "r") as f:
        for line in f:
            x, y, cluster = map(float, line.strip().split(","))
            data.append([x, y, int(cluster)])
    return np.array(data)


# 读取数据
true_centers = read_data("true_centers.txt")
initial_data = read_data("initial_data.txt")
cluster_data = read_data("clustered_data.txt")

# 设置颜色映射
cmap = plt.get_cmap("tab10")
colors = [
    cmap(i) for i in range(len(np.unique(true_centers[:, 2])))
]  # 使用唯一的簇标识数量

# 创建图像
plt.figure(figsize=(15, 6))
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))

# 绘制原始数据
for cluster_id in np.unique(initial_data[:, 2]):
    # 绘制该簇的所有点
    mask = initial_data[:, 2] == cluster_id
    ax1.scatter(
        initial_data[mask][:, 0],
        initial_data[mask][:, 1],
        c=[colors[int(cluster_id)]],
        label=f"Cluster {int(cluster_id)}",
        alpha=0.6,
    )
    # 绘制真实中心点
    center_mask = true_centers[:, 2] == cluster_id
    ax1.scatter(
        true_centers[center_mask][0, 0],
        true_centers[center_mask][0, 1],
        c=[colors[int(cluster_id)]],
        marker="*",
        s=200,
        edgecolor="black",
    )

ax1.set_title("original data distribute")
ax1.legend(bbox_to_anchor=(1.05, 1), loc="upper left")

# 绘制最终聚类结果
for cluster_id in np.unique(cluster_data[:, 2]):
    mask = cluster_data[:, 2] == cluster_id
    ax2.scatter(
        cluster_data[mask][:, 0],
        cluster_data[mask][:, 1],
        c=[colors[int(cluster_id)]],
        label=f"Cluster {int(cluster_id)}",
        alpha=0.6,
    )
    # 计算并绘制最终中心点
    if len(cluster_data[mask]) > 0:
        center_x = np.mean(cluster_data[mask][:, 0])
        center_y = np.mean(cluster_data[mask][:, 1])
        ax2.scatter(
            center_x,
            center_y,
            c=[colors[int(cluster_id)]],
            marker="*",
            s=200,
            edgecolor="black",
        )

ax2.set_title("after K-means")
ax2.legend(bbox_to_anchor=(1.05, 1), loc="upper left")

# 调整布局并保存
plt.tight_layout()
plt.savefig("kmeans_original_vs_clustered.png", bbox_inches="tight", dpi=300)
plt.close()
